{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bf8e01e7-df86-450d-8f01-6573320a93b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyscf import gto, scf, cc, lib\n",
    "import numpy as np\n",
    "from functools import partial\n",
    "import os\n",
    "\n",
    "np.einsum = partial(np.einsum, optimize=True)\n",
    "np.set_printoptions(6, suppress=True, linewidth=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "758eb78f-632c-4f32-82b9-712687fc8d8f",
   "metadata": {},
   "source": [
    "## System Definition and Reference CC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "47412686-c36d-4e0d-bca6-2be749ef60f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "mol = gto.Mole(atom=\"O; H 1 0.94; H 1 0.94 2 104.5\", basis=\"cc-pVDZ\").build()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "623e5681-fa5c-40dd-a44b-8974bfba0f23",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "converged SCF energy = -76.0280410408039\n"
     ]
    }
   ],
   "source": [
    "# this is not the correct way to perform RI-HF with RI-CCSD, these two methods should use different auxbasis\n",
    "# however, energy difference is actually small\n",
    "# and we only want to demonstrate correctness and efficiency of program, so which auxbasis is not really matters\n",
    "mf = scf.RHF(mol).density_fit(auxbasis=\"cc-pVDZ-ri\").run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cc1062e1-0373-4c20-8e9d-49dffe11af24",
   "metadata": {},
   "outputs": [],
   "source": [
    "mol.verbose = 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c49827cb-8a1d-42c6-a969-456d8983f67d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "******** <class 'pyscf.cc.dfccsd.RCCSD'> ********\n",
      "CC2 = 0\n",
      "CCSD nocc = 5, nmo = 24\n",
      "max_cycle = 50\n",
      "direct = 0\n",
      "conv_tol = 1e-07\n",
      "conv_tol_normt = 1e-05\n",
      "diis_space = 6\n",
      "diis_start_cycle = 0\n",
      "diis_start_energy_diff = 1e+09\n",
      "max_memory 4000 MB (current use 158 MB)\n",
      "Init t2, MP2 energy = -76.2307949519314  E_corr(MP2) -0.20275391112758\n",
      "Init E_corr(RCCSD) = -0.20275391112864\n",
      "cycle = 1  E_corr(RCCSD) = -0.207993731903545  dE = -0.00523982077  norm(t1,t2) = 0.0209855\n",
      "cycle = 2  E_corr(RCCSD) = -0.211098740172299  dE = -0.00310500827  norm(t1,t2) = 0.00713562\n",
      "cycle = 3  E_corr(RCCSD) = -0.212120511921666  dE = -0.00102177175  norm(t1,t2) = 0.00266794\n",
      "cycle = 4  E_corr(RCCSD) = -0.21223592969047  dE = -0.000115417769  norm(t1,t2) = 0.000469925\n",
      "cycle = 5  E_corr(RCCSD) = -0.212219794258673  dE = 1.61354318e-05  norm(t1,t2) = 0.000154335\n",
      "cycle = 6  E_corr(RCCSD) = -0.21221718589628  dE = 2.60836239e-06  norm(t1,t2) = 4.34129e-05\n",
      "cycle = 7  E_corr(RCCSD) = -0.212218507545593  dE = -1.32164931e-06  norm(t1,t2) = 1.36829e-05\n",
      "cycle = 8  E_corr(RCCSD) = -0.21221799205749  dE = 5.15488103e-07  norm(t1,t2) = 2.60655e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<class 'pyscf.cc.dfccsd.RCCSD'> does not have attributes  converged\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cycle = 9  E_corr(RCCSD) = -0.212217982473886  dE = 9.58360419e-09  norm(t1,t2) = 6.94985e-07\n",
      "RCCSD converged\n",
      "E(RCCSD) = -76.24025902327774  E_corr = -0.2122179824738857\n"
     ]
    }
   ],
   "source": [
    "mf_cc = cc.CCSD(mf).run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ee93f1f-9950-4f37-8595-92a539fa2842",
   "metadata": {},
   "source": [
    "## Global Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4c70ae44-6571-4917-a36d-e755350acedb",
   "metadata": {},
   "outputs": [],
   "source": [
    "cderi_ao = mf.with_df._cderi\n",
    "mo_occ = mf.mo_occ\n",
    "mo_energy = mf.mo_energy\n",
    "mo_coeff = mf.mo_coeff\n",
    "\n",
    "nmo = len(mo_occ)\n",
    "nocc = (mo_occ > 0).sum()\n",
    "nvir = nmo - nocc\n",
    "nao = mol.nao\n",
    "naux = cderi_ao.shape[0]\n",
    "\n",
    "so = slice(0, nocc)\n",
    "sv = slice(nocc, nmo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "45dd19e3-3201-46f7-990b-cb1012dc60cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "B = np.einsum(\"Puv, up, vq -> pqP\", lib.unpack_tril(cderi_ao), mo_coeff, mo_coeff)\n",
    "D_ov = mo_energy[so, None] - mo_energy[None, sv]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18730ba2-3579-4899-bc51-7a691e3bae83",
   "metadata": {},
   "source": [
    "## Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "73704d45-dc62-4356-a3ef-3c53ae6b7ace",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_intermediates_1(intermediates, B, t1, t2):\n",
    "    # M1j = np.einsum(\"jbP, jb -> P\", B[so, sv], t1)\n",
    "    M1j = t1.reshape(-1) @ B[so, sv].reshape([-1, naux])\n",
    "    \n",
    "    # M1oo = np.einsum(\"jaP, ia -> ijP\", B[so, sv], t1)\n",
    "    M1oo = np.empty([nocc, nocc, naux])\n",
    "    for j in range(nocc):  # par\n",
    "        M1oo[:, j] = t1 @ B[j, sv]\n",
    "    \n",
    "    # M2a = np.einsum(\"jbP, ijab -> iaP\", B[so, sv], (2 * t2 - t2.swapaxes(-1, -2)))\n",
    "    M2a = np.empty([nocc, nvir, naux])\n",
    "    for i in range(nocc):  # par\n",
    "        scr_jba = - t2[i] + 2 * t2.swapaxes(-1, -2)[i]\n",
    "        M2a[i] = scr_jba.reshape([-1, nvir]).T @ B[so, sv].reshape([-1, naux])\n",
    "\n",
    "    intermediates.update({\n",
    "        \"M1j\": M1j,\n",
    "        \"M1oo\": M1oo,\n",
    "        \"M2a\": M2a,\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8aa688b9-7e50-4159-a7e4-d705b9d3e7e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_intermediates_2(intermediates, B, t1, t2):\n",
    "    M1oo = intermediates[\"M1oo\"]\n",
    "    \n",
    "    # M1aov = np.einsum(\"ijP, ja -> iaP\", B[so, so], t1)\n",
    "    M1aov = np.empty([nocc, nvir, naux])\n",
    "    for i in range(nocc):  # par\n",
    "        M1aov[i] = t1.T @ B[i, so]\n",
    "    \n",
    "    # M1bov  = np.einsum(\"abP, ib -> iaP\", B[sv, sv], t1)\n",
    "    M1bov = (t1 @ B[sv, sv].reshape([nvir, -1])).reshape([nocc, nvir, naux])\n",
    "    \n",
    "    # M1vv = np.einsum(\"ibP, ia -> abP\", B[so, sv], t1)\n",
    "    M1vv = (t1.T @ B[so, sv].reshape([nocc, -1])).reshape([nvir, nvir, naux])\n",
    "    \n",
    "    # M2b = np.einsum(\"ikP, ka -> iaP\", M1oo, t1)\n",
    "    M2b = np.empty([nocc, nvir, naux])\n",
    "    for i in range(nocc):\n",
    "        M2b[i] = t1.T @ M1oo[i]\n",
    "\n",
    "    intermediates.update({\n",
    "        \"M1aov\": M1aov,\n",
    "        \"M1bov\": M1bov,\n",
    "        \"M1vv\": M1vv,\n",
    "        \"M2b\": M2b,\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8dc3dcce-2502-4c89-a651-aa6c4a227754",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_residue_1(RHS1, intermediates, B, t1, t2):\n",
    "\n",
    "    M1j = intermediates[\"M1j\"]\n",
    "    M1aov = intermediates[\"M1aov\"]\n",
    "    M1bov = intermediates[\"M1bov\"]\n",
    "    M1oo = intermediates[\"M1oo\"]\n",
    "    M1vv = intermediates[\"M1vv\"]\n",
    "    M2a = intermediates[\"M2a\"]\n",
    "    M2b = intermediates[\"M2b\"]\n",
    "    \n",
    "    # === TERM1 === #\n",
    "    # RHS1 += - 1 * np.einsum(\"lcP, lkP, ikac -> ia\", B[so, sv], M1oo, (2 * t2 - t2.swapaxes(-1, -2)))\n",
    "    \n",
    "    # \"lcP, lkP -> kc\", B[so, sv], M1oo\n",
    "    scr_kc = np.zeros([nocc, nvir])\n",
    "    for l in range(nocc):\n",
    "        scr_kc += M1oo[l] @ B[l, sv].T\n",
    "    \n",
    "    # \"kc, ikac -> ia\", scr_kc, (2 * t2 - t2.swapaxes(-1, -2)))\n",
    "    for i in range(nocc):\n",
    "        t2_i = 2 * t2[i] - t2[i].swapaxes(-1, -2)\n",
    "        t2_i *= scr_kc[:, None, :]\n",
    "        RHS1[i] -= t2_i.sum(axis=(0, 2))\n",
    "    \n",
    "    # === TERM2 === #\n",
    "    # RHS1 += - 1 * np.einsum(\"kcP, icP, ka -> ia\", B[so, sv], (M2a - M1aov), t1)\n",
    "    \n",
    "    RHS1 -= (M2a - M1aov).reshape([nocc, -1]) @ B[so, sv].reshape([nocc, -1]).T @ t1\n",
    "    \n",
    "    # === TERM3 === #\n",
    "    # RHS1 +=   1 * np.einsum(\"icP, acP -> ia\", M2a, B[sv, sv])\n",
    "    \n",
    "    RHS1 += M2a.reshape([nocc, -1]) @ B[sv, sv].reshape([nvir, -1]).T\n",
    "    \n",
    "    # === TERM4 === #\n",
    "    # RHS1 += - 1 * np.einsum(\"ikP, kaP -> ia\", (B[so, so] + M1oo), (M2a + M1bov))\n",
    "    \n",
    "    for k in range(nocc):\n",
    "        RHS1 -= (B[so, k] + M1oo[:, k]) @ (M2a[k] + M1bov[k]).T\n",
    "    \n",
    "    # === TERM5 === #\n",
    "    # RHS1 +=   2 * np.einsum(\"iaP, P -> ia\", (B[so, sv] + M1bov - M1aov + M2a - 0.5 * M2b), M1j)\n",
    "    \n",
    "    scr_iaP = B[so, sv] + M1bov - M1aov + M2a - 0.5 * M2b\n",
    "    RHS1 += 2 * (scr_iaP.reshape([-1, naux]) @ M1j).reshape([nocc, nvir])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5e2a6375-f621-4d99-93a9-aebce6804fc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_residue_2_lt2_contract(RHS2, intermediates, B, t1, t2):\n",
    "\n",
    "    M1j = intermediates[\"M1j\"]\n",
    "    M1aov = intermediates[\"M1aov\"]\n",
    "    M1bov = intermediates[\"M1bov\"]\n",
    "    M1oo = intermediates[\"M1oo\"]\n",
    "    M1vv = intermediates[\"M1vv\"]\n",
    "    M2a = intermediates[\"M2a\"]\n",
    "    M2b = intermediates[\"M2b\"]\n",
    "    \n",
    "    # ====== Loo, Lvv ====== #\n",
    "    \n",
    "    Loo = M2a.reshape([nocc, -1]) @ B[so, sv].reshape([nocc, -1]).T\n",
    "    Loo += ((2 * B[so, so] + M1oo).reshape([-1, naux]) @ M1j).reshape([nocc, nocc])\n",
    "    for l in range(nocc):  # non-par\n",
    "        Loo -= B[l, so] @ M1oo[l].T\n",
    "    \n",
    "    # RHS2 -= np.einsum(\"ik, kjab -> ijab\", Loo, t2)\n",
    "    RHS2 -= (Loo @ t2.reshape([nocc, -1])).reshape([nocc, nocc, nvir, nvir])\n",
    "    \n",
    "    Lvv = ((2 * B[sv, sv] - M1vv).reshape([-1, naux]) @ M1j).reshape([nvir, nvir])\n",
    "    for k in range(nocc):  # non-par\n",
    "        Lvv -= (M2a + M1bov)[k] @ B[k, sv].T\n",
    "    \n",
    "    # RHS2 += np.einsum(\"bc, ijac -> ijab\", Lvv, t2)\n",
    "    RHS2 += (t2.reshape([-1, nvir]) @ Lvv.T).reshape([nocc, nocc, nvir, nvir])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "82656221-a65e-409a-a659-2deb7ea34fcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_residue_2_direct_dot(RHS2, intermediates, B, t1, t2):\n",
    "\n",
    "    M1j = intermediates[\"M1j\"]\n",
    "    M1aov = intermediates[\"M1aov\"]\n",
    "    M1bov = intermediates[\"M1bov\"]\n",
    "    M1oo = intermediates[\"M1oo\"]\n",
    "    M1vv = intermediates[\"M1vv\"]\n",
    "    M2a = intermediates[\"M2a\"]\n",
    "    M2b = intermediates[\"M2b\"]\n",
    "    \n",
    "    # ====== Direct Dot ====== #\n",
    "    \n",
    "    # RHS2 += np.einsum(\"iaP, jbP -> ijab\", B[so, sv] + M2a, 0.5 * B[so, sv] + 0.5 * M2a - M1aov + M1bov - M2b)\n",
    "    # RHS2 -= np.einsum(\"iaP, jbP -> ijab\", M1aov, M1bov)\n",
    "    \n",
    "    scr_iaP = B[so, sv] + M2a\n",
    "    scr_jbP = 0.5 * B[so, sv] + 0.5 * M2a - M1aov + M1bov - M2b\n",
    "    \n",
    "    for i in range(nocc):  # par\n",
    "        for j in range(nocc):  # par\n",
    "            RHS2[i, j] += scr_iaP[i] @ scr_jbP[j].T\n",
    "            RHS2[i, j] -= M1aov[i] @ M1bov[j].T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0c48905e-badc-433a-b61f-f08f34881944",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_residue_2_o3v3(RHS2, intermediates, B, t1, t2):\n",
    "\n",
    "    M1j = intermediates[\"M1j\"]\n",
    "    M1aov = intermediates[\"M1aov\"]\n",
    "    M1bov = intermediates[\"M1bov\"]\n",
    "    M1oo = intermediates[\"M1oo\"]\n",
    "    M1vv = intermediates[\"M1vv\"]\n",
    "    M2a = intermediates[\"M2a\"]\n",
    "    M2b = intermediates[\"M2b\"]\n",
    "    \n",
    "    # ====== O3V3 ====== #\n",
    "    \n",
    "    # t2t = np.ascontiguousarray(t2.swapaxes(-1, -2))\n",
    "    # scr1 = np.einsum(\"kdP, lcP -> kldc\", B[so, sv], B[so, sv])\n",
    "    # scr2 = np.einsum(\"kldc, ilda -> ikca\", scr1, t2)\n",
    "    # scr3 = np.einsum(\"kldc, ilda -> ikca\", scr1, t2t)\n",
    "    # scr4 = np.einsum(\"ikP, acP -> ikca\", (B[so, so] + M1oo), (B[sv, sv] - M1vv))\n",
    "    # RHS2 += np.einsum(\"ikca, jkcb -> ijab\", - scr4 + scr2 - scr3, t2t)\n",
    "    # RHS2 += np.einsum(\"ikcb, jkca -> ijab\", - scr4 + 0.5 * scr2, t2)\n",
    "    \n",
    "    scr_ikP = (B[so, so] + M1oo)\n",
    "    scr_acP = (B[sv, sv] - M1vv)\n",
    "    t2t = np.ascontiguousarray(t2.swapaxes(-1, -2))\n",
    "    \n",
    "    scr1 = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    for k in range(nocc):\n",
    "        for l in range(k + 1):\n",
    "            scr = B[k, sv] @ B[l, sv].T\n",
    "            scr1[k, l] = scr\n",
    "            if k != l:\n",
    "                scr1[l, k] = scr.T\n",
    "    \n",
    "    for i in range(nocc):\n",
    "        scr2 = np.zeros([nocc, nvir, nvir])\n",
    "        scr3 = np.zeros([nocc, nvir, nvir])\n",
    "        for k in range(nocc):\n",
    "            scr2[k] = scr1[k].reshape([-1, nvir]).T @ t2[i].reshape([-1, nvir])\n",
    "            scr3[k] = scr1[k].reshape([-1, nvir]).T @ t2t[i].reshape([-1, nvir])\n",
    "        scr4 = np.zeros([nocc, nvir, nvir])\n",
    "        for c in range(nvir):\n",
    "            scr4[:, c] = scr_ikP[i] @ scr_acP[:, c].T\n",
    "        scr5 = - scr3 - scr4 + scr2\n",
    "        for j in range(nocc):  # par\n",
    "            RHS2[i, j] += scr5.reshape([-1, nvir]).T @ t2t[j].reshape([-1, nvir])\n",
    "        scr5 = - scr4 + 0.5 * scr2\n",
    "        for j in range(nocc):  # par\n",
    "            RHS2[i, j] += t2[j].reshape([-1, nvir]).T @ scr5.reshape([-1, nvir])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "4d8a329a-d1ce-4063-8103-185e92caf119",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_residue_2_o4v2(RHS2, intermediates, B, t1, t2):\n",
    "\n",
    "    M1j = intermediates[\"M1j\"]\n",
    "    M1aov = intermediates[\"M1aov\"]\n",
    "    M1bov = intermediates[\"M1bov\"]\n",
    "    M1oo = intermediates[\"M1oo\"]\n",
    "    M1vv = intermediates[\"M1vv\"]\n",
    "    M2a = intermediates[\"M2a\"]\n",
    "    M2b = intermediates[\"M2b\"]\n",
    "    \n",
    "    # ====== O4V2 (HHL) ====== #\n",
    "    \n",
    "    # Woooo = 0\n",
    "    # Woooo += np.einsum(\"ikP, jlP         -> ijkl\", B[so, so] + M1oo, B[so, so] + M1oo)\n",
    "    # Woooo += np.einsum(\"kcP, ldP, ijcd   -> ijkl\", B[so, sv], B[so, sv], t2)\n",
    "    # RHS2 += 0.5 * np.einsum(\"ijkl, klab   -> ijab\", Woooo, t2,   )\n",
    "    # RHS2 += 0.5 * np.einsum(\"ijkl, ka, lb -> ijab\", Woooo, t1, t1)\n",
    "    \n",
    "    scr_BM1oo = B[so, so] + M1oo\n",
    "    \n",
    "    scr_ijkl = np.zeros([nocc, nocc, nocc, nocc])\n",
    "    for i in range(nocc):\n",
    "        for j in range(i + 1):\n",
    "            tmp_ab = scr_BM1oo[i] @ scr_BM1oo[j].T\n",
    "            scr_ijkl[i, j] = tmp_ab\n",
    "            scr_ijkl[j, i] = tmp_ab.T\n",
    "    scr_klcd = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    for k in range(nocc):\n",
    "        for l in range(k + 1):\n",
    "            tmp_cd = B[k, sv] @ B[l, sv].T\n",
    "            scr_klcd[k, l] = tmp_cd\n",
    "            scr_klcd[l, k] = tmp_cd.T\n",
    "    scr_ijkl += (t2.reshape([nocc * nocc, -1]) @ scr_klcd.reshape([nocc * nocc, -1]).T).reshape([nocc, nocc, nocc, nocc])\n",
    "    \n",
    "    tau2 = t2 + t1[:, None, :, None] * t1[None, :, None, :]\n",
    "    RHS2 += 0.5 * (scr_ijkl.reshape([nocc * nocc, -1]) @ tau2.reshape([-1, nvir * nvir])).reshape(t2.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "75994061-e3d8-44b2-921e-b3b46e015613",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_residue_2_o2v4(RHS2, intermediates, B, t1, t2):\n",
    "\n",
    "    M1j = intermediates[\"M1j\"]\n",
    "    M1aov = intermediates[\"M1aov\"]\n",
    "    M1bov = intermediates[\"M1bov\"]\n",
    "    M1oo = intermediates[\"M1oo\"]\n",
    "    M1vv = intermediates[\"M1vv\"]\n",
    "    M2a = intermediates[\"M2a\"]\n",
    "    M2b = intermediates[\"M2b\"]\n",
    "    \n",
    "    # ====== O2V4 (PPL) ====== #\n",
    "    \n",
    "    # Wvvvv = 0\n",
    "    # Wvvvv += np.einsum(\"acP, bdP -> abcd\", B[sv, sv] - M1vv, B[sv, sv] - M1vv)\n",
    "    # Wvvvv -= np.einsum(\"acP, bdP -> abcd\", M1vv, M1vv)\n",
    "    # RHS2 += 0.5 * np.einsum(\"abcd, ijcd   -> ijab\", Wvvvv, t2,   )\n",
    "    # RHS2 += 0.5 * np.einsum(\"abcd, ic, jd -> ijab\", Wvvvv, t1, t1)\n",
    "    \n",
    "    # scr_acP = (B[sv, sv] - M1vv)\n",
    "    \n",
    "    # for a in range(nvir):\n",
    "    #     for b in range(a + 1):\n",
    "    #         tmp_cd = scr_acP[a] @ scr_acP[b].T\n",
    "    #         tmp_cd -= M1vv[a] @ M1vv[b].T\n",
    "    #         tmp_ij = 0.5 * tau2.reshape([nocc, nocc, nvir * nvir]) @ tmp_cd.reshape(-1)\n",
    "    #         RHS2[:, :, a, b] += tmp_ij\n",
    "    #         if a != b:\n",
    "    #             RHS2[:, :, b, a] += tmp_ij.T\n",
    "    \n",
    "    RHS_PPL = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    \n",
    "    scr_acP = (B[sv, sv] - M1vv)\n",
    "    tau2 = t2 + t1[:, None, :, None] * t1[None, :, None, :]\n",
    "    \n",
    "    batch_a = 8\n",
    "    batch_b = 32\n",
    "    for a_start in range(0, nvir, batch_a):  # par\n",
    "        a_end = min(a_start + batch_a, nvir)\n",
    "        nbatch_a = a_end - a_start\n",
    "        sa = slice(a_start, a_end)\n",
    "        for b_start in range(0, a_end, batch_b):  # par\n",
    "            b_end = min(b_start + batch_b, a_end)\n",
    "            nbatch_b = b_end - b_start\n",
    "            sb = slice(b_start, b_end)\n",
    "    \n",
    "            scr_abcd = np.zeros([nbatch_a, nbatch_b, nvir, nvir])\n",
    "            for a in range(nbatch_a):  # par\n",
    "                for b in range(nbatch_b):  # par\n",
    "                    scr_abcd[a, b] += scr_acP[a + a_start] @ scr_acP[b + b_start].T\n",
    "                    scr_abcd[a, b] -= M1vv[a + a_start] @ M1vv[b + b_start].T\n",
    "            scr_ijab = 0.5 * tau2.reshape([nocc * nocc, nvir * nvir]) @ scr_abcd.reshape([nbatch_a * nbatch_b, nvir * nvir]).T\n",
    "            scr_ijab = scr_ijab.reshape([nocc, nocc, nbatch_a, nbatch_b])\n",
    "            RHS_PPL[:, :, sa, sb] = scr_ijab\n",
    "    for i in range(nocc):  # par\n",
    "        for j in range(nocc):  # par\n",
    "            for a in range(nvir):\n",
    "                for b in range(a):\n",
    "                    RHS_PPL[i, j, b, a] = RHS_PPL[j, i, a, b]\n",
    "    \n",
    "    RHS2 += RHS_PPL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9b883746-9ad6-40f0-9234-5b2c03c96bdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_energy(intermediates, B, t1, t2):\n",
    "    M1j = intermediates[\"M1j\"]\n",
    "    M1oo = intermediates[\"M1oo\"]\n",
    "    M2a = intermediates[\"M2a\"]\n",
    "\n",
    "    e_t1 = 2 * M1j.reshape(-1) @ M1j.reshape(-1)\n",
    "    e_t1 -= M1oo.reshape(-1) @ M1oo.swapaxes(0, 1).reshape(-1)\n",
    "    e_t2 = B[so, sv].reshape(-1) @ M2a.reshape(-1)\n",
    "    return e_t1, e_t2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0e5228be-1e28-4810-ab71-b9dbccbe586f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_amplitude(intermediates, B, mo_energy, e_corr, t1, t2):\n",
    "    get_riccsd_intermediates_2(intermediates, B, mf_cc.t1, mf_cc.t2)\n",
    "\n",
    "    RHS1 = np.zeros([nocc, nvir])\n",
    "    get_riccsd_residue_1(RHS1, intermediates, B, mf_cc.t1, mf_cc.t2)\n",
    "\n",
    "    RHS2 = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    get_riccsd_residue_2_lt2_contract(RHS2, intermediates, B, mf_cc.t1, mf_cc.t2)\n",
    "    get_riccsd_residue_2_direct_dot(RHS2, intermediates, B, mf_cc.t1, mf_cc.t2)\n",
    "    get_riccsd_residue_2_o3v3(RHS2, intermediates, B, mf_cc.t1, mf_cc.t2)\n",
    "    get_riccsd_residue_2_o2v4(RHS2, intermediates, B, mf_cc.t1, mf_cc.t2)\n",
    "    get_riccsd_residue_2_o4v2(RHS2, intermediates, B, mf_cc.t1, mf_cc.t2)\n",
    "    # RHS2 = RHS2 + RHS2.transpose((1, 0, 3, 2))\n",
    "\n",
    "    # update t1\n",
    "    D_ov = mo_energy[so, None] - mo_energy[None, sv]\n",
    "    RHS1 /= D_ov\n",
    "    t1_new = RHS1\n",
    "\n",
    "    # update t2\n",
    "    for i in range(nocc):  # par\n",
    "        for j in range(i + 1):  # par\n",
    "            t2_ab = RHS2[i, j] + RHS2[j, i].T\n",
    "            d2_ab = (D_ov[i][:, None] + D_ov[j][None, :])\n",
    "            t2_ab /= d2_ab\n",
    "            RHS2[i, j] = t2_ab\n",
    "            if i != j:\n",
    "                RHS2[j, i] = t2_ab.T\n",
    "    t2_new = RHS2\n",
    "\n",
    "    delta_t1 = np.abs(t1_new - t1).sum()\n",
    "    delta_t2 = np.abs(t2_new - t2).sum()\n",
    "\n",
    "    get_riccsd_intermediates_1(intermediates, B, t1_new, t2_new)\n",
    "    e1_new, e2_new = get_riccsd_energy(intermediates, B, t1_new, t2_new)\n",
    "    e_corr_new = e1_new + e2_new\n",
    "\n",
    "    delta_e_corr = np.abs(e_corr_new - e_corr)\n",
    "\n",
    "    print(f\"delta_e_corr {delta_e_corr}\")\n",
    "    print(f\"delta_t1     {delta_t1}\")\n",
    "    print(f\"delta_t2     {delta_t2}\")\n",
    "    print(f\"e_corr       {e_corr_new}\")\n",
    "    \n",
    "    return (e_corr_new, t1_new, t2_new), (delta_e_corr, delta_t1, delta_t2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "b30572b5-2ad2-4e52-ba64-41d709fcdcae",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_riccsd_amplitude(intermediates, B, mo_energy, e_corr, t1, t2):\n",
    "    get_riccsd_intermediates_2(intermediates, B, t1, t2)\n",
    "\n",
    "    RHS1 = np.zeros([nocc, nvir])\n",
    "    get_riccsd_residue_1(RHS1, intermediates, B, t1, t2)\n",
    "\n",
    "    RHS2 = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    get_riccsd_residue_2_lt2_contract(RHS2, intermediates, B, t1, t2)\n",
    "    get_riccsd_residue_2_direct_dot(RHS2, intermediates, B, t1, t2)\n",
    "    get_riccsd_residue_2_o3v3(RHS2, intermediates, B, t1, t2)\n",
    "    get_riccsd_residue_2_o2v4(RHS2, intermediates, B, t1, t2)\n",
    "    get_riccsd_residue_2_o4v2(RHS2, intermediates, B, t1, t2)\n",
    "    # RHS2 = RHS2 + RHS2.transpose((1, 0, 3, 2))\n",
    "\n",
    "    # update t1\n",
    "    D_ov = mo_energy[so, None] - mo_energy[None, sv]\n",
    "    RHS1 /= D_ov\n",
    "    t1_new = RHS1\n",
    "\n",
    "    # update t2\n",
    "    for i in range(nocc):  # par\n",
    "        for j in range(i + 1):  # par\n",
    "            t2_ab = RHS2[i, j] + RHS2[j, i].T\n",
    "            d2_ab = (D_ov[i][:, None] + D_ov[j][None, :])\n",
    "            t2_ab /= d2_ab\n",
    "            RHS2[i, j] = t2_ab\n",
    "            if i != j:\n",
    "                RHS2[j, i] = t2_ab.T\n",
    "    t2_new = RHS2\n",
    "\n",
    "    delta_t1 = np.abs(t1_new - t1).sum()\n",
    "    delta_t2 = np.abs(t2_new - t2).sum()\n",
    "\n",
    "    get_riccsd_intermediates_1(intermediates, B, t1_new, t2_new)\n",
    "    e1_new, e2_new = get_riccsd_energy(intermediates, B, t1_new, t2_new)\n",
    "    e_corr_new = e1_new + e2_new\n",
    "\n",
    "    delta_e_corr = np.abs(e_corr_new - e_corr)\n",
    "\n",
    "    print(f\"delta_e_corr {delta_e_corr}\")\n",
    "    print(f\"delta_t1     {delta_t1}\")\n",
    "    print(f\"delta_t2     {delta_t2}\")\n",
    "    print(f\"e_corr       {e_corr_new}\")\n",
    "    \n",
    "    return (e_corr_new, t1_new, t2_new), (delta_e_corr, delta_t1, delta_t2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72a571b7-bc10-43dc-8c2e-dddd186eff8a",
   "metadata": {},
   "source": [
    "## Actual CCSD Iterations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7458431a-57e4-45f5-b94e-223fe7a35164",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "delta_e_corr 0.21221786522503241\n",
      "delta_t1     1.5344182097907008e-06\n",
      "delta_t2     8.552378380389298e-06\n",
      "e_corr       -0.21221786522503241\n"
     ]
    }
   ],
   "source": [
    "intermediates = {\n",
    "    \"M1j\": None,\n",
    "    \"M1oo\": None,\n",
    "    \"M1aov\": None,\n",
    "    \"M1bov\": None,\n",
    "    \"M1vv\": None,\n",
    "    \"M2a\": None,\n",
    "    \"M2b\": None,\n",
    "}\n",
    "\n",
    "get_riccsd_intermediates_1(intermediates, B, mf_cc.t1, mf_cc.t2)\n",
    "get_riccsd_amplitude(intermediates, B, mo_energy, 0, mf_cc.t1, mf_cc.t2);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "df8ce573-1a92-4a69-85f2-cebd531f5069",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MP2 e_corr -0.2027539177355094\n",
      "=== iteration 0 ===\n",
      "delta_e_corr 0.005239822061174626\n",
      "delta_t1     0.04577406620457161\n",
      "delta_t2     0.6288253059568694\n",
      "e_corr       -0.20799373979668404\n",
      "=== iteration 1 ===\n",
      "delta_e_corr 0.0031049880275015596\n",
      "delta_t1     0.00899702551326303\n",
      "delta_t2     0.19035492947826982\n",
      "e_corr       -0.2110987278241856\n",
      "=== iteration 2 ===\n",
      "delta_e_corr 0.0006920906765417456\n",
      "delta_t1     0.00442466132512682\n",
      "delta_t2     0.06188439259412255\n",
      "e_corr       -0.21179081850072734\n",
      "=== iteration 3 ===\n",
      "delta_e_corr 0.0002674380057528236\n",
      "delta_t1     0.0016531799635276863\n",
      "delta_t2     0.02408455765575785\n",
      "e_corr       -0.21205825650648016\n",
      "=== iteration 4 ===\n",
      "delta_e_corr 9.53670364023107e-05\n",
      "delta_t1     0.0007457361270473629\n",
      "delta_t2     0.009718102230343598\n",
      "e_corr       -0.21215362354288247\n",
      "=== iteration 5 ===\n",
      "delta_e_corr 3.803422412995694e-05\n",
      "delta_t1     0.00033367856551322724\n",
      "delta_t2     0.004124376139367448\n",
      "e_corr       -0.21219165776701243\n",
      "=== iteration 6 ===\n",
      "delta_e_corr 1.521383044489344e-05\n",
      "delta_t1     0.00015773357709401362\n",
      "delta_t2     0.001799964388037765\n",
      "e_corr       -0.21220687159745732\n",
      "=== iteration 7 ===\n",
      "delta_e_corr 6.313098083604585e-06\n",
      "delta_t1     7.658198946943465e-05\n",
      "delta_t2     0.000815379459031079\n",
      "e_corr       -0.21221318469554093\n",
      "=== iteration 8 ===\n",
      "delta_e_corr 2.65009137387362e-06\n",
      "delta_t1     3.880323438192166e-05\n",
      "delta_t2     0.0003824015199805647\n",
      "e_corr       -0.2122158347869148\n",
      "=== iteration 9 ===\n",
      "delta_e_corr 1.1335458867389736e-06\n",
      "delta_t1     2.13975682140287e-05\n",
      "delta_t2     0.00018580448111444058\n",
      "e_corr       -0.21221696833280154\n",
      "=== iteration 10 ===\n",
      "delta_e_corr 4.911838537569047e-07\n",
      "delta_t1     1.2215420913371917e-05\n",
      "delta_t2     9.311618421346301e-05\n",
      "e_corr       -0.2122174595166553\n",
      "=== iteration 11 ===\n",
      "delta_e_corr 2.1609474154526964e-07\n",
      "delta_t1     7.107639208112146e-06\n",
      "delta_t2     4.8370340456718196e-05\n",
      "e_corr       -0.21221767561139684\n",
      "=== iteration 12 ===\n",
      "delta_e_corr 9.648018567398609e-08\n",
      "delta_t1     4.127776212511925e-06\n",
      "delta_t2     2.5858635167092027e-05\n",
      "e_corr       -0.21221777209158252\n",
      "=== iteration 13 ===\n",
      "delta_e_corr 4.378568424501239e-08\n",
      "delta_t1     2.393967379721769e-06\n",
      "delta_t2     1.4096019340706168e-05\n",
      "e_corr       -0.21221781587726676\n",
      "=== iteration 14 ===\n",
      "delta_e_corr 2.0221434970935093e-08\n",
      "delta_t1     1.386946333405018e-06\n",
      "delta_t2     7.824164508415365e-06\n",
      "e_corr       -0.21221783609870173\n"
     ]
    }
   ],
   "source": [
    "t1 = np.zeros([nocc, nvir])\n",
    "t2 = np.zeros([nocc, nocc, nvir, nvir])\n",
    "e_corr = 0\n",
    "for i in range(nocc):\n",
    "    for j in range(i + 1):\n",
    "        d2_ab = (D_ov[i][:, None] + D_ov[j][None, :])\n",
    "        t2_ab = (B[i, sv] @ B[j, sv].T) / d2_ab\n",
    "        t2[i, j] = t2_ab\n",
    "        factor = 2 if i != j else 1\n",
    "        e_corr += factor * (t2_ab * (2 * t2_ab - t2_ab.T) * d2_ab).sum()\n",
    "        if i != j:\n",
    "            t2[j, i] = t2_ab.T\n",
    "\n",
    "print(f\"MP2 e_corr {e_corr}\")\n",
    "\n",
    "intermediates = {\n",
    "    \"M1j\": None,\n",
    "    \"M1oo\": None,\n",
    "    \"M1aov\": None,\n",
    "    \"M1bov\": None,\n",
    "    \"M1vv\": None,\n",
    "    \"M2a\": None,\n",
    "    \"M2b\": None,\n",
    "}\n",
    "\n",
    "for i in range(100):\n",
    "    print(f\"=== iteration {i} ===\")\n",
    "    get_riccsd_intermediates_1(intermediates, B, t1, t2)\n",
    "    (e_corr, t1, t2), (delta_e_corr, delta_t1, delta_t2) = get_riccsd_amplitude(intermediates, B, mo_energy, e_corr, t1, t2)\n",
    "    if delta_e_corr < 1e-7 and delta_t1 < 1e-5 and delta_t2 < 1e-5:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b21a916-05a0-4d79-af93-0f4e2ab4c141",
   "metadata": {},
   "source": [
    "## Save temporaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b8d35cb5-5328-4fd3-9710-a09a34f3fec6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_riccsd_temporaries(intermediates, B, mo_energy):\n",
    "    t1 = mf_cc.t1\n",
    "    t2 = mf_cc.t2\n",
    "\n",
    "    os.makedirs(\"h2o-cc-pvdz\", exist_ok=True)\n",
    "    np.save(\"h2o-cc-pvdz/mo_coeff.npy\", np.ascontiguousarray(mf.mo_coeff))\n",
    "    np.save(\"h2o-cc-pvdz/mo_occ.npy\", mf.mo_occ)\n",
    "    np.save(\"h2o-cc-pvdz/mo_energy.npy\", mf.mo_energy)\n",
    "    np.save(\"h2o-cc-pvdz/t1.npy\", t1)\n",
    "    np.save(\"h2o-cc-pvdz/t2.npy\", t2)\n",
    "    np.save(\"h2o-cc-pvdz/cderi.npy\", mf.with_df._cderi)\n",
    "    \n",
    "    get_riccsd_intermediates_1(intermediates, B, t1, t2)\n",
    "    get_riccsd_intermediates_2(intermediates, B, t1, t2)\n",
    "\n",
    "    np.save(\"h2o-cc-pvdz/M1j.npy\", intermediates[\"M1j\"])\n",
    "    np.save(\"h2o-cc-pvdz/M1oo.npy\", intermediates[\"M1oo\"])\n",
    "    np.save(\"h2o-cc-pvdz/M1aov.npy\", intermediates[\"M1aov\"])\n",
    "    np.save(\"h2o-cc-pvdz/M1bov.npy\", intermediates[\"M1bov\"])\n",
    "    np.save(\"h2o-cc-pvdz/M1vv.npy\", intermediates[\"M1vv\"])\n",
    "    np.save(\"h2o-cc-pvdz/M2a.npy\", intermediates[\"M2a\"])\n",
    "    np.save(\"h2o-cc-pvdz/M2b.npy\", intermediates[\"M2b\"])\n",
    "\n",
    "    RHS1 = np.zeros([nocc, nvir])\n",
    "    get_riccsd_residue_1(RHS1, intermediates, B, t1, t2)\n",
    "    np.save(\"h2o-cc-pvdz/RHS1.npy\", RHS1)\n",
    "\n",
    "    RHS2 = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    RHS2_tmp = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    get_riccsd_residue_2_lt2_contract(RHS2_tmp, intermediates, B, t1, t2)\n",
    "    RHS2 += RHS2_tmp\n",
    "    np.save(\"h2o-cc-pvdz/RHS2-lt2_contract.npy\", RHS2_tmp)\n",
    "    \n",
    "    RHS2_tmp = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    get_riccsd_residue_2_direct_dot(RHS2_tmp, intermediates, B, t1, t2)\n",
    "    RHS2 += RHS2_tmp\n",
    "    np.save(\"h2o-cc-pvdz/RHS2-direct_dot.npy\", RHS2_tmp)\n",
    "    \n",
    "    RHS2_tmp = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    get_riccsd_residue_2_o3v3(RHS2_tmp, intermediates, B, t1, t2)\n",
    "    RHS2 += RHS2_tmp\n",
    "    np.save(\"h2o-cc-pvdz/RHS2-o3v3.npy\", RHS2_tmp)\n",
    "    \n",
    "    RHS2_tmp = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    get_riccsd_residue_2_o2v4(RHS2_tmp, intermediates, B, t1, t2)\n",
    "    RHS2 += RHS2_tmp\n",
    "    np.save(\"h2o-cc-pvdz/RHS2-o2v4.npy\", RHS2_tmp)\n",
    "    \n",
    "    RHS2_tmp = np.zeros([nocc, nocc, nvir, nvir])\n",
    "    get_riccsd_residue_2_o4v2(RHS2_tmp, intermediates, B, t1, t2)\n",
    "    RHS2 += RHS2_tmp\n",
    "    np.save(\"h2o-cc-pvdz/RHS2-o4v2.npy\", RHS2_tmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "077d65d8-0e1a-4e79-8ee1-44105072293b",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_riccsd_temporaries(intermediates, B, mo_energy)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
